"""
Normalization
"""
import numpy as np

from mindspore import Tensor
from mindspore.common import dtype as mstype


class Normalization:
    """
    数据归一化处理类
    用于对湍流模型的输入特征和输出标签进行标准化处理
    支持MinMax归一化和标准差归一化两种方法
    """
    def __init__(self, feature_method, label_method, file_path):
        self.feature_method = feature_method  # 特征归一化方法
        self.label_method = label_method      # 标签归一化方法
        self.file_path = file_path            # 标准化参数文件路径

    def normalize_feature(self, data):
        """
        对输入特征进行归一化处理
        
        Args:
            data: 输入特征数据
            
        Returns:
            Tensor: 归一化后的特征数据
        """
        if self.feature_method == "MinMax":
            return self.feature_min_max_norm(data)
        if self.feature_method == "Std":
            return self.feature_std_norm(data)
        return self.feature_min_max_norm(data)  # 默认使用MinMax归一化

    def normalize_label(self, mut_value, uy_value, vx_value):
        """
        对输出标签及相关物理量进行归一化处理
        
        Args:
            mut_value: 涡粘性系数
            uy_value: u对y的偏导数(∂u/∂y)
            vx_value: v对x的偏导数(∂v/∂x)
            
        Returns:
            tuple: 归一化后的涡粘性系数、应变率和雷诺应力
        """
        if self.label_method == "MinMax":
            return self.label_min_max_norm(mut_value, uy_value, vx_value)
        if self.label_method == "Std":
            return self.label_std_norm(mut_value, uy_value, vx_value)
        return self.label_min_max_norm(mut_value, uy_value, vx_value)  # 默认使用MinMax归一化

    def feature_min_max_norm(self, data):
        '''
        特征数据的MinMax归一化
        公式: (x - min)/(max - min)
        
        Args:
            data: 原始特征数据
            
        Returns:
            Tensor: MinMax归一化后的特征数据
        '''
        df_max, df_min = get_min_max_data_from_txt(self.file_path)
        fenmu = (df_max[:, :-1] - df_min[:, :-1])  # 归一化分母（最大值-最小值）
        fenzi = (data - df_min[:, :-1])            # 归一化分子（数据-最小值）
        df_data = fenzi / fenmu                    # 归一化计算
        return df_data

    def feature_std_norm(self, data):
        '''
        特征数据的标准差标准化
        公式: (x - mean)/std
        
        Args:
            data: 原始特征数据
            
        Returns:
            Tensor: 标准差标准化后的特征数据
        '''
        df_mean, df_std = get_mean_std_data_from_txt(self.file_path)
        fenmu = df_std[:, :-1]                     # 标准化分母（标准差）
        fenzi = (data - df_mean[:, :-1])           # 标准化分子（数据-均值）
        df_data = fenzi / fenmu                    # 标准化计算
        return df_data

    def label_min_max_norm(self, mut_value, uy_value, vx_value):
        '''
        标签数据的MinMax归一化（包含物理量计算）
        对涡粘性、应变率和雷诺应力进行归一化处理
        
        Args:
            mut_value: 涡粘性系数
            uy_value: u对y的偏导数(∂u/∂y)
            vx_value: v对x的偏导数(∂v/∂x)
            
        Returns:
            tuple: 归一化后的涡粘性、应变率和雷诺应力
        '''
        df_max, df_min = get_min_max_data_from_txt(self.file_path)
        # 涡粘性系数MinMax归一化
        norm_mut = (mut_value - df_min[:, -1]) / (df_max[:, -1] - df_min[:, -1])
        
        # 计算应变率并归一化
        sij = 0.5 * (uy_value + vx_value) / 6500  # 应变率计算并缩放到合理范围
        sij = (sij - sij.min()) / (sij.max() - sij.min())  # MinMax归一化
        
        # 计算雷诺应力（基于Boussinesq假设）
        rs_value = 2 * sij * norm_mut  # 雷诺应力 = 2 * 涡粘性 * 应变率
        
        return norm_mut, sij, rs_value

    def label_std_norm(self, mut_value, uy_value, vx_value):
        '''
        标签数据的标准差标准化（包含物理量计算）
        对涡粘性、应变率和雷诺应力进行标准化处理
        
        Args:
            mut_value: 涡粘性系数
            uy_value: u对y的偏导数(∂u/∂y)
            vx_value: v对x的偏导数(∂v/∂x)
            
        Returns:
            tuple: 标准化后的涡粘性、雷诺应力和应变率
        '''
        df_mean, df_std = get_mean_std_data_from_txt(self.file_path)
        # 涡粘性系数标准差标准化
        norm_mut = (mut_value - df_mean[:, -1]) / df_std[:, -1]
        
        # 计算应变率并标准化
        sij = 0.5 * (uy_value + vx_value) / 6500  # 应变率计算并缩放到合理范围
        sij = (sij - sij.mean()) / sij.std()  # 标准差标准化
        
        # 计算雷诺应力（基于Boussinesq假设）
        rs_value = 2 * sij * norm_mut  # 雷诺应力 = 2 * 涡粘性 * 应变率
        
        return norm_mut, rs_value, sij


def get_min_max_data_from_txt(data_path):
    """
    从文本文件读取MinMax归一化参数（最大值和最小值）
    
    Args:
        data_path: 数据路径
        
    Returns:
        tuple: 最大值和最小值的Tensor
    """
    df_max = np.loadtxt(data_path + '/2d_max.txt')
    df_min = np.loadtxt(data_path + '/2d_min.txt')
    df_max = Tensor(df_max, mstype.float32).reshape(1, -1)
    df_min = Tensor(df_min, mstype.float32).reshape(1, -1)
    return df_max, df_min


def get_mean_std_data_from_txt(data_path):
    """
    从文本文件读取标准差标准化参数（均值和标准差）
    
    Args:
        data_path: 数据路径
        
    Returns:
        tuple: 均值和标准差的Tensor
    """
    df_mean = np.loadtxt(data_path + '/mean.txt')
    df_std = np.loadtxt(data_path + '/std.txt')
    df_mean = Tensor(df_mean, mstype.float32).reshape(1, -1)
    df_std = Tensor(df_std, mstype.float32).reshape(1, -1)
    return df_mean, df_std